#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import sys
import argparse
import itertools
from math import isinf, isnan
from pprint import pprint as pp 
import my
try:
    import MySQLdb
    import MySQLdb.cursors
except ImportError:
    raise ImportError,"The Python MySQLdb module is required to run this program. Try 'pip install MySQL-python'."

def AddPathwayAnnotationsError(): pass

#-i /scratch1/tmp/myourshaw/rnaseq_20140227/results/gene_exp.diff.head

def is_int(s):
    try:
        n = int(s)
        if (isinf(n) or isnan(n)):
            return False
        return True
    except ValueError:
        return False


def run(input, output=None,
        skip_string=None, ignore_lines=None, header_line='#',
        column_names=None, no_column_names=False,
        host='cortex.local', port=3306, user='vax', password='vax',
        database='vax_test', table='pathways_concat_subpathways',
        gene_col='gene',ensg_col=None,):

    if not output:
        output=input+'.pathways'
    header_line_number = int(header_line) if is_int(header_line) and int(header_line) > 0 else None
    header_chars = header_line if not is_int(header_line) else None    
    table = '`{}`.`{}`'.format(database,table)
    if ensg_col:
        input_key = ensg_col #in input file, e.g., test_id
        table_key = 'ENSG'
    else:
        input_key = gene_col #in input file, e.g., gene
        table_key = 'gene'


    print('Cheking database connection')
    with MySQLdb.connect(host,user,password,port=port) as cursor:
        try:
            query = 'SELECT VERSION()'
            cursor.execute(query)
            mysql_version = cursor.fetchone()
        except MySQLdb.Error, e:
            msg = ''
            try:
                msg = 'MySQL Error [{}]: {}'.format(e.args[0], e.args[1])
            except IndexError:
                msg = 'MySQL Error {}'.format(e)
            raise AddPathwayAnnotationsError('Database connection failed. {}'.format(msg))
        else:
            if not mysql_version:
                raise AddPathwayAnnotationsError('Database connected but SELECT VERSION() query failed.')
            print('Database connected. MySQL version: {}'.format(mysql_version))
            query = "SELECT SCHEMA_NAME FROM INFORMATION_SCHEMA.SCHEMATA WHERE SCHEMA_NAME = '{}'".format(database)
            cursor.execute(query)
            schema = cursor.fetchone()
            print('Database `{}` {}.'.format(database, "doesn't exist" if schema == None else "exists"))
            query = 'SELECT DISTINCT pathway FROM {}'.format(table)
            cursor.execute(query)
            pathways = sorted(list(itertools.chain.from_iterable(cursor.fetchall())))
    with open(input,'r') as input_fh, open(output,'w') as output_fh:
        print('Annotating {}.'.format(input))
        column_names_list = None
        line_count = 0
        first_data_line = None
        header_found = False if header_line_number or header_chars else True
        header_written = False
        if column_names:
            header_found = True
            column_names_list = column_names
            column_names = {i:column_names_list[i] for i in range(len(column_names_list))}
        else:
            column_names = {}
        for line in input_fh:
            line_count += 1
            line = line.rstrip('\r\n')
            if not bool(line.strip()):
                continue
            if not header_found:
                if no_column_names:
                    header_found = True
                    column_names_list = ['col'+str(i+1) for i in range(len(line.split('\t')))]
                    column_names = {i:column_names_list[i] for i in range(len(column_names_list))}
                elif ((header_line_number and header_line_number == line_count)
                      or (header_chars and header_chars != '#' and line.startswith(header_chars))
                      or (header_chars and header_chars == '#' and line.startswith(header_chars) and line[1] != '#')):
                    header_found = True
                    column_names_list = line.split('\t')
                continue
            elif (skip_string and line.startswith(skip_string)) or (ignore_lines and line_count <= ignore_lines):
                continue
            else: #data line
                if first_data_line == None:
                    first_data_line = line_count
                if not header_written:
                    header = column_names_list + pathways
                    header_dict = {header[i]: i for i in range(len(header))}
                    output_fh.write('{}\n'.format('\t'.join(header)))
                    header_written = True
                fields = line.split('\t')
                output_dict = [i for i in range(len(fields))]			
                table_value = fields[header_dict[input_key]]
                pathway_dict = {p: '.' for p in pathways}
                query = """SELECT `pathway`, `sub_pathway` FROM {table} WHERE `{key}` = '{value}'""".format(
                    table=table, key=table_key, value=table_value)
                with MySQLdb.connect(host=host, user=user, passwd=password, port=port, db=database, cursorclass=MySQLdb.cursors.DictCursor) as cursor:
                    try:
                        cursor.execute(query)
                        data = cursor.fetchall()
                    except MySQLdb.Error, e:
                        msg = ''
                        try:
                            msg = 'MySQL Error [{}]: {}'.format(e.args[0], e.args[1])
                        except IndexError:
                            msg = 'MySQL Error {}'.format(e)
                        raise AddPathwayAnnotationsError('Database connection failed. {}'.format(msg))
                    except:
                        data = []
                if data != None:
                    for d in data:
                        pathway_dict[d['pathway']] = d['sub_pathway']
                pathway_output = [pathway_dict[p] for p in pathways]
                output_data = fields + pathway_output
                output_fh.write('{}\n'.format('\t'.join(output_data)))
    print('Done.')



def main():
    #command line arguments
    parser = argparse.ArgumentParser(
        description = 'add pathway annotations to a file',
        epilog = """1.0β1 ©2014 Michael Yourshaw all rights reserved.""")
    parser.add_argument('--input', '-i', type=str, metavar='FILE', required=True,
                        help='Tab-delimited input file')
    parser.add_argument('--output', '-o', type=str, metavar='FILE',
                        help='Ouput file (default: <input>.pathways')
    parser.add_argument('--skip_string',
                        help='skip lines that start with this string')
    parser.add_argument('--ignore_lines', type=int,
                        help='ignore the first n lines')
    parser.add_argument('--header_line', default = '#',
                        help='1-based row number of column names or # if the first row that starts with a single # character')
    parser.add_argument('--column_names', nargs='*',
                        help='list of column names (when file does not contain column name header)')
    parser.add_argument('--no_column_names', action='store_true', default=False,
                        help='file does not contain column name header, so use "col1, col2, ..."')
    parser.add_argument('--host', '-H', type=str, default='cortex.local',
                        help='MySQL database server hostname or ip address (default: cortex.local)')
    parser.add_argument('--port', '-P', type=int, default=3306,
                        help='MySQL database (default: 3306)')
    parser.add_argument('--user', '-u', type=str, default='vax',
                        help='MySQL database user (default: vax)')
    parser.add_argument('--password', '-p', type=str, default='vax',
                        help='MySQL password (default: None = enter silently at prompt)')
    parser.add_argument('--database', '-d', type=str, default='vax_test',
                        help='MySQL database (default: vax_test)')
    parser.add_argument('--table', '-t', type=str, default='pathways_concat_subpathways',
                        help='MySQL table (default: vax_test)')
    parser.add_argument('--gene_col', '-g', type=str, default='gene',
                        help='Column name of HGNC gene symbol in input file (default: gene)')
    parser.add_argument('--ensg_col', '-e', type=str,
                        help='Column name of ENSG in input file; overrides --gene (default: None)')
    args = parser.parse_args()    

    run(input=args.input, output=args.output,
        skip_string=args.skip_string, ignore_lines=args.ignore_lines, header_line=args.header_line,
        column_names=args.column_names, no_column_names=args.no_column_names,
        host=args.host, port=args.port, user=args.user, password=args.password,
        database=args.database, table=args.table, gene_col=args.gene_col, ensg_col=args.ensg_col,)

if __name__ == "__main__": sys.exit(main())